from datetime import datetime
from time import mktime

import numpy as np
import tables as tb

import matplotlib.pyplot as plt

# Some helper functions
conv = {0: lambda s: mktime(datetime.strptime(s, "%Y-%m-%d").timetuple())}

def load_timeseries(path):
    """Loads finacial time series data into memory as a numpy array."""
    arr = np.loadtxt(path, skiprows=1, delimiter=',', converters=conv)
    arr = arr[::-1]
    return arr

# Load all of the time series data
GE  = load_timeseries('GE.csv')
GM  = load_timeseries('GM.csv')
IBM = load_timeseries('IBM.csv')
JNJ = load_timeseries('JNJ.csv')
MMM = load_timeseries('MMM.csv')


# Make some more meaningful indices.
DATE = 0
OPEN = 1
HIGH = 2
LOW = 3
CLOSE = 4
VOLUME = 5
ADJ_CLOSE = 6

# <demo> stop

# Make new pytables file
h5 = tb.openFile('financial_data.h5', 'w')

# <demo> stop

# Data in HDF5 is stored in a heirarchy 
# that resembles the Unix filesystem.
# Groups are anaglogous to directories.
# Datasets are similar to files, and like 
# files, there are a lot of different types.
# The standard ones are Arrays and Tables.

# First, like the Unix filesystem, we have a fundemental root directory
h5r = h5.root   # Pythonic way

h5r = h5.getNode('/')   # filesystem way

# <demo> stop


# However, root is not the most useful place to store things, 
# Let's create another directory.
adj_close_group = h5.createGroup(h5r, 'adj_close', 'Adjusted Close')

h5.flush()
# <demo> stop

# Now, let's add arrays that represent the different stocks' adjusted closing values to this group.

h5.createArray('/adj_close', 'GE', GE[:, ADJ_CLOSE])

gm_array = h5.createArray('/adj_close', 'GM', GM[:, ADJ_CLOSE])

ibm_array = h5.createArray(adj_close_group, 'IBM', IBM[:, ADJ_CLOSE])

h5.createArray(adj_close_group, 'JNJ', JNJ[:, ADJ_CLOSE], title='Johnson & Johnson Adjusted Close')

h5.createArray(adj_close_group, 'MMM', MMM, title='3M Adjusted Close')

h5.flush()
# <demo> stop

# Finally, Create a time array that applies over all data
time_array = h5.createArray(h5.root, 'time', GE[:, DATE])

h5.flush()
# <demo> stop



# Now that we have these arrays we can access and change their values 
print "The first ten entries for IBM are:", 
print ibm_array[:10]
print

h5.flush()
# <demo> stop

# Now, let's say GM is worthless
gm_array[:] = 0.0

h5.flush()
# <demo> stop

# Since, GM is worthless, you may as well get rid of it.
h5.removeNode('/adj_close/GM')

h5.flush()
# <demo> stop

# However, life is not simply full of arrays. 
# The concept of tables, or structured arrays, is also very useful.
# In their 2D expression, these are similar to SQL tables.
# Here, data fields are aranged into columns, while data entries are rows.

# Let's create another group to store the raw ticker data as tables...
ticker_group = h5.createGroup('/', 'ticker', 'Raw Ticker Data')

h5.flush()
# <demo> stop

# To create a Table, we need to describe the fields, or columns.
ticker_desc = {'date': tb.Time64Col(pos=DATE), 
               'open': tb.Float64Col(pos=OPEN), 
               'high': tb.Float64Col(pos=HIGH), 
               'low':  tb.Float64Col(pos=LOW), 
               'close': tb.Float64Col(pos=CLOSE), 
               'volume': tb.Float64Col(pos=VOLUME, dflt=42.0), 
               'adj_close': tb.Float64Col(pos=ADJ_CLOSE), 
               }

# Other columns types include:
#   tb.Float32Col()
#   tb.Int8Col() 
#   tb.StringCol(itemsize=20) 

# <demo> stop

# Now, we shall add the raw data to the database
ge_table = h5.createTable('/ticker', 'GE', ticker_desc)
ge_table.append(GE)

h5.flush()
# <demo> stop

# Alternatively, you can loop through the data, one row at a time.
h5.createTable('/ticker', 'GM', ticker_desc)

# Note, that as long as the table has a 'natural name' (ie, a name that 
# could also be a python variable), you can get a reference to a node
# in the heirarch, using dot-access notation.
gm_table = h5.root.ticker.GM

# continue with adding data to table
h5_row = gm_table.row

for gm_row in GM:
    h5_row['date'] = gm_row[DATE]
    h5_row['open'] = gm_row[OPEN]
    h5_row['high'] = gm_row[HIGH]
    h5_row['low'] = gm_row[LOW]
    h5_row['close'] = gm_row[CLOSE]
    h5_row['adj_close'] = gm_row[ADJ_CLOSE]

    # You have to append this row on each loop through!
    h5_row.append()

h5.flush()
# <demo> stop

# Finally, Let's say that you have a portfolio that 
# starts out with everything invested in General Electric 
# and slowly moves to the present, where everything is 
# invested in IBM.  A bold move!

h5.createGroup('/adj_close', 'portfolio', 'My Stock Portfolio Data')

prop = np.linspace(0.0, 1.0, len(time_array))
my_portfolio = (h5r.adj_close.GE[:] * prop[::-1]) + (ibm_array[:] * prop)

h5.createArray('/adj_close/portfolio', 'mine', my_portfolio)

h5.flush()
# <demo> stop

# Now, to visualize this portfolio
plt.plot(time_array, ibm_array, 'b-', label='IBM')
plt.plot(time_array, h5r.adj_close.GE, 'r-', label='GE')
plt.plot(time_array, h5r.adj_close.portfolio.mine, 'g-', label='My Portfolio')

plt.legend(loc=0)
plt.xlabel('Time since Epoch [s]')
plt.ylabel('Adjusted Closing Value [$]')

plt.show()

# Always remember to close your file when you are done with it!
h5.close()
# <demo> stop
 
